package com.zzyy.study.controller;

import cn.hutool.core.date.DateUtil;
import com.zzyy.study.service.AiAssistant;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.output.Response;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @auther zzyy
 * @create 2025-03-08 11:38
 */
@RestController
@Slf4j
public class ChatMemoryController
{
    @Resource
    private ChatLanguageModel chatLanguageModel;
    private ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(10);

    @Resource
    private AiAssistant aiAssistant;

    private Map<Long,ChatMemory> chatMemoryMap = new ConcurrentHashMap<>();

    @GetMapping(value = "/chatmemory/lowapi")
    public String lowapi(@RequestParam("prompt") String prompt)
    {
        chatMemory.add(UserMessage.userMessage(prompt));
        Response<AiMessage> aiMessageResponse = chatLanguageModel.generate(chatMemory.messages());
        chatMemory.add(aiMessageResponse.content());

        return aiMessageResponse.content().text();
    }

    @GetMapping(value = "/chatmemory/highapi")
    public String highapi(@RequestParam("userId") Long userId,@RequestParam("prompt") String prompt)
    {
        String chat = aiAssistant.chat(userId, prompt);

        return "testWithChatMemory success : "+ DateUtil.now()+"\t"+chat;
    }

    @GetMapping(value = "/chatmemory/highapi/manual")
    public String manual(@RequestParam("userId") Long userId,@RequestParam("prompt") String prompt)
    {
        if(!chatMemoryMap.containsKey(userId))
        {
            MessageWindowChatMemory messageWindowChatMemory = MessageWindowChatMemory.withMaxMessages(10);
            messageWindowChatMemory.add(UserMessage.userMessage(prompt));
            chatMemoryMap.put(userId,messageWindowChatMemory);
        }
        List<ChatMessage> messages = chatMemoryMap.get(userId).messages();
        Response<AiMessage> aiMessageResponse = chatLanguageModel.generate(messages);
        chatMemoryMap.get(userId).add(aiMessageResponse.content());

        return aiMessageResponse.content().text();

    }
}
